import sys

sys.path.append("./lib")
import rosbag

from matplotlib.pyplot import cla
from sensor_msgs import point_cloud2
import cv2

import os
import json
import numpy as np
from tqdm import tqdm
import argparse
import threading
# from cv_bridge import CvBridge, CvBridgeError
#
# bridge = CvBridge()

from .utils.ddd_utils import draw_box_3d, project_to_image, compute_box_3d, rot_y2alpha, alpha2rot_y

needed_choice = ['left_front', 'front_middle', 'right_front', 'left_back', 'right_back']

Calibs = {
    "left_front": {
        "intrinsic":
            [9.8136791215576147e+02, 0., 9.4839461394621117e+02, 0., 0.,
             9.7877930144468655e+02, 5.4826911056422762e+02, 0., 0., 0., 1., 0., ],
        "extrinsic":
            [-0.489303, -0.872083, -0.00725849, 0.316294,
             -0.0682373, 0.0465809, -0.996582, -0.122966,
             0.86944, -0.487136, -0.0823008, 0.174502,
             0, 0, 0, 1]
    },

    "front_middle": {
        "intrinsic":
            [9.6329358074361460e+02, 0., 1.0077748249221447e+03, 0., 0.,
             9.6501237970612237e+02, 5.4370143226399443e+02, 0., 0., 0., 1., 0., ],
        "extrinsic":
            [-0.00325009, 0.999501, -0.0314101, 0.0161996,
             -0.972325, 0.0041792, 0.233596, -1.90735,
             0.233611, 0.0313, 0.971827, -0.156316,
             0, 0, 0, 1]
    },

    "right_front": {
        "intrinsic":
            [9.9156983385556464e+02, 0., 9.5979848329359470e+02, 0., 0.,
             9.9052320215775444e+02, 5.3216083204490576e+02, 0., 0., 0., 1., 0.],
        "extrinsic":
            [0.38748, -0.920667, 0.0472163, -0.516601,
             -0.00512604, -0.0533686, -0.998562, -0.0777872,
             0.921863, 0.386681, -0.0253985, 0.684171,
             0, 0, 0, 1]
    },

    "left_back": {
        "intrinsic":
            [9.7623060203681041e+02, 0., 9.4639844986342735e+02, 0., 0.,
             9.7682136279895110e+02, 5.4500167321113429e+02, 0., 0., 0., 1., 0., ],
        "extrinsic":
            [0.991594, 0.129371, -0.00216614, 0.0266785,
             -0.0018602, -0.00248605, -0.999995, 0.339616,
             -0.129376, 0.991594, -0.00222448, -0.0266944,
             0, 0, 0, 1]
    },

    "right_back": {
        "intrinsic":
            [9.8506879793285611e+02, 0., 1.0038607766204225e+03, 0., 0.,
             9.8454037729996605e+02, 5.4638714145710173e+02, 0., 0., 0., 1., 0., ],
        "extrinsic":
            [-0.977639, 0.205631, 0.0440235, 0.292283,
             -0.0489388, -0.0188763, -0.998624, -0.206619,
             -0.204517, -0.978447, 0.0285176, 0.316011,
             0, 0, 0, 1]
    },

}
#  {'/iv_points':"iv",'/rslidar_points':"m1",'/os_cloud_node_1/points':"ouster1",'/os_cloud_node_2/points':"ouster2"}
cam2lidar_map = {'left_front': '/os_cloud_node_1/points', 'front_middle': '/iv_points',
                 'right_front': "/os_cloud_node_2/points", 'left_back': '/os_cloud_node_1/points',
                 'right_back': "/os_cloud_node_2/points"}


def loadlabel(label_path):
    with open(label_path) as f:
        labels = json.load(f)
        objs = []
        for label in labels:
            obj = {}
            position = label['psr']['position']
            rotation = label['psr']['rotation']['z']
            scale = label['psr']['scale']
            x, y, z = position["x"], position["y"], position["z"]
            xs, ys, zs = scale["x"], scale["y"], scale["z"]
            obj["obj_type"], obj["obj_id"] = label['obj_type'], label['obj_id']
            obj["box"] = [x, y, z, xs, ys, zs, rotation]
            objs.append(obj)
    return objs


def write_pcd(points, save_pcd_path):
    # points = []
    # for p in gen:
    #     points.append(p)
    n = len(points)
    lines = []
    # print("n",n)
    for i in range(n):
        # x, y, z, i, is_g, data_time = points[i]
        x, y, z, i = points[i][:4]
        lines.append('{:.6f} {:.6f} {:.6f} {}'.format( \
            x, y, z, i))
    with open(save_pcd_path, 'w') as f:
        f.write(HEADER.format(n, n))
        f.write('\n'.join(lines))


def save_lidar(path, pair_dict):
    points = []
    for topic in arg.base_lidar:
        lidar_dict = pair_dict[topic]
        gen = lidar_dict["data"]
        for p in gen:
            points.append(p)
    write_pcd(points, path)


def pair_to_B(A_list, B_list):
    A_pair = []
    B_pair = []
    # 进行配对
    for idx, A in enumerate(A_list):
        time_gap_min = 100
        A_value = A[0]
        A_time = A[1]
        for B in B_list:
            if abs(A_time - B[1]) < time_gap_min:
                time_gap_min = abs(A_time - B[1])
                B_value = B[0]
                B_time = B[1]

        A_pair.append([A_value, A_time])
        if time_gap_min > arg.time_threshold:
            B_pair.append([None, None])
        else:
            B_pair.append([B_value, B_time])
    return [A_pair, B_pair]


def process_main(bag_path):
    # 加载raw数据
    print("=========process start!===========")
    res_bag_list = []
    bag_file_list = os.listdir(bag_path)
    labels = get_all_labels()
    totalnum = len(bag_file_list)
    t = tqdm(enumerate(bag_file_list))
    for bag_idx, bag_file in t:
        # if bag_idx <=10:
        #     continue
        # if not "2022-09-05-16-08-58_7" in bag_file:
        #     continue
        t.set_description("[{}/{}]".format(bag_idx, totalnum))
        bag_file_path = os.path.join(bag_path, bag_file)
        if not os.path.isfile(bag_file_path):
            continue
        [bag_name, file_type] = bag_file.split("/")[-1].split(".")
        if not file_type == "bag":
            continue
        # global data_dict
        data_dict = process_bag(bag_file_path)
        # return
        # compare_msg(data)
        # if arg.thread:
        #     threading.Thread(target=precess_,args=(bag_idx, bag_file)).start()
        # else:
        #     precess_(bag_idx, bag_file)
        #     print("===================over===================")
        #     break
        # break
        process_cam(data_dict, labels)

    print("=========process over!===========")


def get_label(boxs_3d):
    # bboxs_3d.append([dim, location, rotation_y, box_2d[-1],class_name])
    # col1 ： 类标签
    # col5-8 ：2D bbox，左上和右下坐标
    # col9-11 : 3D bbox dim 高、宽、长
    # col12-14 ：3D bbox centor 坐标
    # col15: 车体方向角
    labels = []
    for box_3d in boxs_3d:
        # class_name = class_map[box_3d[4]]
        class_name = box_3d[4]
        obj_id = box_3d[5]
        b3 = box_3d[3]
        [bbox_xy0, bbox_xy1] = [max(b3[:, 0]), max(b3[:, 1])], [min(b3[:, 0]), min(b3[:, 1])]
        dim = box_3d[0]
        location = [_[0] for _ in box_3d[1].tolist()]
        rotation = box_3d[2]
        # label = [[class_name,class_id],0,0,0,bbox_xy0[0],bbox_xy0[1],bbox_xy1[0],bbox_xy1[1]]+dim+location+[rotation]
        label = [class_name, obj_id, bbox_xy0[0], bbox_xy0[1], bbox_xy1[0], bbox_xy1[1]] + dim + location + [rotation]
        labels.append(label)
    return labels


def cam2lidar(img, orientation, pair_dict, cam_file_path, labels):
    imgc = img.copy()
    lidarname = cam2lidar_map[orientation]

    bin_name = pair_dict[lidarname]["bin_name"]
    # print(bin_name,orientation)
    josn_name = bin_name.replace(".bin", ".json")
    if josn_name not in labels:
        return
    label_path = os.path.join(labels[josn_name], josn_name)
    objs = loadlabel(label_path)
    imgc = img.copy()
    labels_list, bird_view = showddd(imgc, Calibs[orientation], objs, orientation)
    if (bird_view is None):
        return
    cv2.imwrite(cam_file_path, img)
    cv2.imwrite(cam_file_path.replace("/images", "/images_3dbbox"), imgc)
    cv2.imwrite(cam_file_path.replace("/images", "/images_bev").format(orientation), bird_view)

    calib_file_path = cam_file_path.replace("/images", "/calibs").replace(".png", ".json")
    with open(calib_file_path, 'w') as f:
        json.dump(Calibs[orientation], f, indent=4, )

    label_file_path = cam_file_path.replace("/images", "/labels").replace(".png", ".txt")
    save_label(labels_list, label_file_path)

    # import matplotlib.pyplot as plt
    # from PIL import Image
    # imgc = Image.fromarray(cv2.cvtColor(imgc, cv2.COLOR_BGR2RGB))
    # plt.imshow(imgc)

    def process_bag(bag_file):
        print("=============>load {}".format(bag_file.split("/")[-1]))
        bag = rosbag.Bag(bag_file)
        info = bag.get_type_and_topic_info()
        class_dict = info.topics
        topics = arg.base_lidar + arg.camera_topic
        bag_data = bag.read_messages(topics)
        data_dict = {}
        time_list = []
        data_dict["bag_name"] = bag_file.split("/")[-1].split(".")[0]
        for t in topics:
            data_dict[t] = []

        for idx, (topic, msg, t) in enumerate(bag_data):
            time = msg.header.stamp.secs + msg.header.stamp.nsecs / 1000000000.0
            if topic in arg.base_lidar:
                msg = point_cloud2.read_points(msg)
            data_dict[topic].append({"time": time, "data": msg})
            # time_list.append({"time":time,"topic":topic,"data":msg})
        return data_dict


topic2name = {'/iv_points': "iv", '/rslidar_points': "m1", '/os_cloud_node_1/points': "ouster1",
              '/os_cloud_node_2/points': "ouster2"}


# arg.ignore_class=[]
def idx2binname(idx, bagname, topic):
    idx_str = str(idx).zfill(4)
    lidar_name = topic2name[topic]
    bin_name = "{}_frame_{}_{}.bin".format(bagname, lidar_name, idx_str)
    return bin_name


def find_pair_lidar(cam_time, data_dict):
    pair_dict = {}
    for topic in arg.base_lidar:
        lidar_dict = data_dict[topic]
        time_gap_min = 10000
        best_data = None
        best_idx = -1
        for idx, dd in enumerate(lidar_dict):
            time_gap = abs(dd["time"] - cam_time)
            if (time_gap < 1) and (time_gap < time_gap_min):
                time_gap_min = time_gap
                best_idx = idx + 1
                best_data = dd["data"]
        bin_name = idx2binname(best_idx, data_dict["bag_name"], topic)
        # pair_dict[topic]={"bin_name":bin_name,"data":best_data}
        pair_dict[topic] = {"bin_name": bin_name}
    return pair_dict


def get_all_labels():
    p1 = "/data/percetion/boleidun/"
    labels = {}
    for pi in os.listdir(p1):
        p2 = os.path.join(p1, pi, "label")
        if os.path.isdir(p2):
            for pii in os.listdir(p2):
                labels[pii] = p2
    return labels


def process_cam(data_dict, labels):
    pair_map = {}
    cam_topic = arg.camera_topic[0]
    orientation_choice = ['Top_view', 'left_front', 'front_middle', 'right_front', 'left_back', 'right_back',
                          'rearview',
                          'Infrared']
    needed_choice = ['left_front', 'front_middle', 'right_front', 'left_back', 'right_back']
    for idx, dd in enumerate(data_dict[cam_topic]):
        cam_time = dd["time"]
        cam_data = dd["data"]

        pair_dict = find_pair_lidar(cam_time, data_dict)
        # bin_path = "./boleidun/lidar/{}_{}_{}.bin".format(data_dict["bag_name"],"lidar",str(idx).zfill(4))
        # save_lidar(bin_path,pair_dict)

        for cam_i, image in enumerate(cam_data.raw_image):
            orientation = orientation_choice[cam_i]
            if (orientation not in needed_choice):
                continue
            cam_file_path = "./boleidun/images/{}_{}_{}.png".format(data_dict["bag_name"], orientation,
                                                                    str(idx).zfill(4))
            img = bridge.imgmsg_to_cv2(image, desired_encoding='bgr8')
            # if arg.distortion:
            #     img = cv2.undistort(img, np.array(Breton_intrinsic[orientation]), np.array(aian_K[orientation]))

            cam2lidar(img, orientation, pair_dict, cam_file_path, labels)
            # if(orientation=="front_middle"):
            # break
            # showddd(img)
        # break

        map_file_path = "./boleidun/maps/{}_{}_{}.txt".format(data_dict["bag_name"], "all", str(idx).zfill(4))
        with open(map_file_path, "w") as f:
            for pd in pair_dict:
                bin_name = pair_dict[pd]["bin_name"]
                f.write(bin_name + "\n")

        # break


args = argparse.ArgumentParser()
args.add_argument("--bags", default="/data/percetion/博雷顿现场数据/0907/2/问题包", type=str, dest='bags')  # 0907/2/问题包
args.add_argument("--target", default="/data/percetion/fuyu/test", type=str)
args.add_argument("--force", default=True)
args.add_argument("--time_threshold", type=float, default=0.05, dest='time_threshold')
args.add_argument("--pair_topic", nargs='+', type=str, default=['cam', 'points'], dest='pair_topic')
args.add_argument("--base_lidar", type=str,
                  default=['/iv_points', '/rslidar_points', '/os_cloud_node_1/points', '/os_cloud_node_2/points'],
                  dest='base_lidar')
args.add_argument("--camera_topic", type=str,
                  default=['/n_camera_gmsl_prep'])
args.add_argument("--distortion", type=bool, default=False, dest='distortion')
args.add_argument("--save_lidar_type", type=str, default='bin', dest='save_lidar_type')
args.add_argument("--thread", action='store_true', default=False, )
args.add_argument("--ignore_class", type=str, default=['Wall', "Dust"])
arg = args.parse_args([])
# arg.thread = True
# dirs = ["images_bev","images_3dbbox","images","calibs","labels","maps"]
# for d in dirs:
#     os.system("rm -r ./boleidun/{} ; mkdir ./boleidun/{}".format(d,d))
process_main(arg.bags)

# 0：iv   /iv_points
# 1：rslidar point /rslidar_points
# 2: ouster left /os_cloud_node_1/points
# 3: ouster right /os_cloud_node_2/points
# camera  /n_camera_gmsl_prep
